#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Dict, Any, Optional, Tuple
import asyncio

from app.core.logging import get_logger, LoggerMixin
from app.core.exceptions import NotFoundError, DatabaseError, EmbeddingError, LLMError, SQLGenerationError
from app.core.utils import sanitize_sql
from app.repositories.training_data import TrainingDataRepository
from app.repositories.vector import VectorRepository
from app.repositories.query_history import QueryHistoryRepository
from app.services.embedding import EmbeddingService
from app.services.llm import LLMManager
from app.models.database import TrainingData, QueryHistory

# 创建日志记录器
logger = get_logger("text2sql_service")


class Text2SQLService(LoggerMixin):
    """文本到SQL服务
    
    管理文本到SQL转换的核心功能
    """
    
    def __init__(
        self,
        training_data_repo: TrainingDataRepository,
        vector_repo: VectorRepository,
        query_history_repo: QueryHistoryRepository,
        embedding_service: EmbeddingService,
        llm_manager: LLMManager
    ):
        """初始化文本到SQL服务
        
        Args:
            training_data_repo: 训练数据仓库
            vector_repo: 向量仓库
            query_history_repo: 查询历史仓库
            embedding_service: 嵌入服务
            llm_manager: LLM管理器
        """
        self.training_data_repo = training_data_repo
        self.vector_repo = vector_repo
        self.query_history_repo = query_history_repo
        self.embedding_service = embedding_service
        self.llm_manager = llm_manager
    
    async def _get_similar_questions(
        self,
        project_id: str,
        question: str,
        top_k: int = 5
    ) -> List[TrainingData]:
        """获取相似问题
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[TrainingData]: 相似问题列表
        """
        try:
            # 生成问题嵌入
            vector = await self.embedding_service.embed_query(question)
            
            # 搜索相似向量，过滤出question_sql类型
            similar_items = await self.vector_repo.search_similar(
                project_id=project_id,
                data_type="train",  # 使用统一的data_type
                vector=vector,
                top_k=top_k,
                content_type="question_sql"  # 只获取question_sql类型的数据
            )
            
            if not similar_items:
                return []
            
            # 获取内容ID列表
            content_ids = [item["content_id"] for item in similar_items]
            
            # 获取问题-SQL数据
            question_sql_list = await self.training_data_repo.get_question_sql_by_ids(content_ids)
            
            return question_sql_list
            
        except Exception as e:
            self.logger.error(f"获取相似问题失败: {str(e)}")
            return []
    
    async def _get_relevant_ddl(
        self,
        project_id: str,
        question: str,
        top_k: int = 3
    ) -> List[TrainingData]:
        """获取相关DDL
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[TrainingData]: 相关DDL列表
        """
        try:
            # 生成问题嵌入
            vector = await self.embedding_service.embed_query(question)
            
            # 搜索相似向量，过滤出ddl类型
            similar_items = await self.vector_repo.search_similar(
                project_id=project_id,
                data_type="train",  # 使用统一的data_type
                vector=vector,
                top_k=top_k,
                content_type="ddl"  # 只获取ddl类型的数据
            )
            
            if not similar_items:
                # 如果没有找到相似的DDL，尝试获取所有DDL
                ddl_list = await self.training_data_repo.get_ddl_by_project(project_id)
                return ddl_list[:top_k]  # 返回最多top_k个DDL
            
            # 获取内容ID列表
            content_ids = [item["content_id"] for item in similar_items]
            
            # 获取DDL数据
            ddl_list = await self.training_data_repo.get_ddl_by_ids(content_ids)
            
            return ddl_list
            
        except Exception as e:
            self.logger.error(f"获取相关DDL失败: {str(e)}")
            return []
    
    async def _get_similar_history(
        self,
        project_id: str,
        question: str,
        top_k: int = 3
    ) -> List[QueryHistory]:
        """获取相似历史记录
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[QueryHistory]: 相似历史记录列表
        """
        try:
            # 简单的文本匹配
            query_history_list = await self.query_history_repo.search_by_question(
                project_id=project_id,
                query=question,
                limit=top_k
            )
            return query_history_list
        except Exception as e:
            self.logger.error(f"获取相似历史记录失败: {str(e)}")
            return []
    
    async def _get_relevant_descriptions(
        self,
        project_id: str,
        question: str,
        top_k: int = 3
    ) -> List[TrainingData]:
        """获取相关描述
        
        Args:
            project_id: 项目ID
            question: 问题
            top_k: 返回的最大结果数
            
        Returns:
            List[TrainingData]: 相关描述列表
        """
        try:
            # 生成问题嵌入
            vector = await self.embedding_service.embed_query(question)
            
            # 搜索相似向量，过滤出description类型
            similar_items = await self.vector_repo.search_similar(
                project_id=project_id,
                data_type="train",  # 使用统一的data_type
                vector=vector,
                top_k=top_k,
                content_type="description"  # 只获取description类型的数据
            )
            
            if not similar_items:
                return []
            
            # 获取内容ID列表
            content_ids = [item["content_id"] for item in similar_items]
            
            # 获取描述数据
            description_list = await self.training_data_repo.get_descriptions_by_ids(content_ids)
            
            return description_list
            
        except Exception as e:
            self.logger.error(f"获取相关描述失败: {str(e)}")
            return []
    
    async def _build_context(
        self,
        project_id: str,
        question: str
    ) -> str:
        """构建上下文
        
        Args:
            project_id: 项目ID
            question: 问题
            
        Returns:
            str: 上下文
        """
        # 并行获取相关数据
        similar_questions_task = self._get_similar_questions(project_id, question)
        relevant_ddl_task = self._get_relevant_ddl(project_id, question)
        similar_history_task = self._get_similar_history(project_id, question)
        relevant_descriptions_task = self._get_relevant_descriptions(project_id, question)
        
        # 等待所有任务完成
        similar_questions, relevant_ddl, similar_history, relevant_descriptions = await asyncio.gather(
            similar_questions_task,
            relevant_ddl_task,
            similar_history_task,
            relevant_descriptions_task
        )
        
        # 构建上下文
        context_parts = []
        
        # 添加DDL
        if relevant_ddl:
            context_parts.append("## 数据库模式")
            for i, ddl in enumerate(relevant_ddl):
                context_parts.append(f"### 表定义 {i+1}")
                context_parts.append(ddl.content)
            context_parts.append("")
        
        # 添加描述信息
        if relevant_descriptions:
            context_parts.append("## 数据库描述")
            for i, desc in enumerate(relevant_descriptions):
                context_parts.append(f"### 描述 {i+1}")
                context_parts.append(desc.content)
            context_parts.append("")
        
        # 添加相似问题
        if similar_questions:
            context_parts.append("## 相似问题和SQL")
            for i, q in enumerate(similar_questions):
                context_parts.append(f"### 示例 {i+1}")
                context_parts.append(f"问题: {q.question}")
                context_parts.append(f"SQL: {q.sql}")
            context_parts.append("")
        
        # 添加历史记录
        if similar_history:
            context_parts.append("## 历史查询")
            for i, h in enumerate(similar_history):
                context_parts.append(f"### 历史 {i+1}")
                context_parts.append(f"问题: {h.question}")
                context_parts.append(f"SQL: {h.generated_sql}")
            context_parts.append("")
        
        # 如果没有上下文，添加提示
        if not context_parts:
            context_parts.append("没有找到相关的数据库模式或示例。请根据问题生成合理的SQL查询。")
        
        return "\n".join(context_parts)
    
    async def generate_sql(
        self,
        project_id: str,
        question: str,
        db_type: str = "mysql",
        system_prompt: Optional[str] = None
    ) -> str:
        """生成SQL查询
        
        Args:
            project_id: 项目ID
            question: 问题
            db_type: 数据库类型
            system_prompt: 自定义系统提示词
            
        Returns:
            str: 生成的SQL查询
            
        Raises:
            SQLGenerationError: 当生成失败时
        """
        try:
            # 构建上下文
            context = await self._build_context(project_id, question)
            
            # 生成SQL
            sql = await self.llm_manager.generate_sql(
                question=question,
                context=context,
                db_type=db_type,
                system_prompt=system_prompt
            )
            
            # 清理SQL
            clean_sql = sanitize_sql(sql)
            
            # 记录历史
            await self.query_history_repo.create(
                project_id=project_id,
                history_type="generation",
                question=question,
                generated_sql=clean_sql,
                db_type=db_type
            )
            
            self.logger.info(f"已生成SQL查询，长度: {len(clean_sql)}")
            return clean_sql
            
        except LLMError as e:
            error_msg = f"生成SQL查询失败: {str(e)}"
            self.logger.error(error_msg)
            raise SQLGenerationError(error_msg)
        except Exception as e:
            error_msg = f"生成SQL查询失败: {str(e)}"
            self.logger.error(error_msg)
            raise SQLGenerationError(error_msg)